package com.lry.basic.algorithm.trie.demo2;

public class Trie {

    private class Node{

        public Node[] nodes;
        public boolean isWord;

        public Node(){
            nodes = new Node[26];
        }


        public Node get(char c){
            return nodes[c-'a'];
        }

        public boolean containsKey(char c){
            return get(c)!=null;
        }

        public void put(char c ,Node node){
            nodes[c-'a'] = node;
        }
    }

    private int size;
    private Node root;
    public Trie(){
        root = new Node();
    }


    public void add(String word){
        char[] cs = checkWord(word);
        Node cur = root;
        for(char c:cs){
            if(!cur.containsKey(c)){
                cur.put(c,new Node());
            }
            cur = cur.get(c);
        }
        if(!cur.isWord){
            cur.isWord = true;
            ++size;
        }
    }

    private char[] checkWord(String word){
        if(null==word||word.isEmpty())
            throw new RuntimeException("word must not null or empty");
        char[]cs = word.toCharArray();

        for(char c:cs){
            if(!(c>='a'&&c<='z')){
                throw new RuntimeException("字符必须a~z");
            }
        }
        return cs;
    }

    private char[] checkPattern(String pattern){
        if(null==pattern||pattern.isEmpty())
            throw new RuntimeException("pattern must not null or empty");
        char[]cs = pattern.toCharArray();

        for(char c:cs){
            if('.'==c)
                continue;
            if(!(c>='a'&&c<='z')){
                throw new RuntimeException("字符必须a~z");
            }
        }
        return cs;
    }

    public boolean contains(String word){
        char[] cs = checkWord(word);
        Node cur = root;
        for(char c:cs){
            if(!cur.containsKey(c))
                return false;
            cur = cur.get(c);
        }
        return cur.isWord;
    }

    //pattern可以写.匹配一个字符
    public boolean match(String pattern){
        char[] cs = checkPattern(pattern);

        return match(root,cs,0);
    }

    private boolean match(Node node, char[] cs, int index) {
        if(null==node)
            return false;

        if(index==cs.length){
            return node.isWord;
        }

        if(cs[index]=='.'){
            for(Node n:node.nodes){
                if(match(n,cs,index+1)){
                    return true;
                }
            }
            return false;
        }else{
            if(!node.containsKey(cs[index])){
                return false;
            }
            return match(node.get(cs[index]),cs,index+1);
        }
    }


    // 查询是否在Trie中有单词以prefix为前缀
    public boolean isPrefix(String prefix){
        char[] cs = checkWord(prefix);
        Node cur = root;
        for(char c:cs){
            if(!cur.containsKey(c))
                return false;
            cur = cur.get(c);
        }
        return true;
    }


    public int size() {
        return size;
    }


    public static void main(String[] args) {
        Trie trie = new Trie();
        trie.add("hello");
        System.out.println(trie.contains("hello"));
        System.out.println(trie.match("he.o"));
        System.out.println(trie.isPrefix("he"));
    }
}
